package njust.csc.aeapi.service

import io.github.oshai.kotlinlogging.KotlinLogging
import kotlinx.coroutines.*
import kotlinx.coroutines.sync.withLock
import njust.csc.aeapi.config.ApplicationVariables
import njust.csc.aeapi.dao.AnalysisOrderRepository
import njust.csc.aeapi.dao.ApkWhitelistRepository
import njust.csc.aeapi.dao.SignatureBlacklistRepository
import njust.csc.aeapi.dao.WordBlacklistRepository
import njust.csc.aeapi.entity.AnalysisOrder
import njust.csc.aeapi.entity.JudgeResult
import njust.csc.aeapi.utils.GlobalMutex
import njust.csc.aecore.analysis.DataProcess
import njust.csc.aecore.analysis.DynamicExplorer
import njust.csc.aecore.analysis.StaticExplorer
import njust.csc.aecore.entity.APK
import njust.csc.aecore.judge.GPTJudge
import njust.csc.aecore.utils.PathUtils
import njust.csc.aecore.utils.SerializeUtils
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Service
import java.nio.file.Path
import java.sql.Timestamp
import kotlin.io.path.*

private val logger = KotlinLogging.logger {}

/**
 * @Author ZhouHang
 * @Date 24/6/19 周三
 */
@Service
class AnalysisService @Autowired constructor(
    private val analysisOrderRepository: AnalysisOrderRepository,
    private val wordBlacklistRepository: WordBlacklistRepository,
    private val signatureBlacklistRepository: SignatureBlacklistRepository,
    private val apkWhitelistRepository: ApkWhitelistRepository,
    private val coroutineScope: CoroutineScope,
    applicationVariables: ApplicationVariables,
) {
    private val gpt = GPTJudge()
    private val dataProcess = DataProcess()
    private val userOutputPath = Path(applicationVariables.userOutputDirPath)
    private val userInputPath = Path(applicationVariables.userInputDirPath)

    /**
     * 分析 apkName 指定的 apk 文件
     * @param apkName 待分析的 apk 文件名
     * @param analysisNo 分析号
     * */
    fun analysis(apkName: String, analysisNo: String) {
        coroutineScope.launch(Dispatchers.Default) {
            val resultPath = userOutputPath.resolve(analysisNo).resolve("result.json")
            val order = getAnalysisOrder(analysisNo)
            if (order == null) {
                // 该 apk 从未分析过
                var analysisOrder = AnalysisOrder()
                analysisOrder.createTime = Timestamp(System.currentTimeMillis())
                analysisOrder.staticStatus = "Analysing"
                analysisOrder.dynamicStatus = "Analysing"
                analysisOrder.judgeStatus = "Judging"
                analysisOrder.judgeResult = "-1"
                analysisOrder.apkName = apkName
                analysisOrder.analysisNo = analysisNo
                analysisOrder = withContext(Dispatchers.IO) {
                    analysisOrderRepository.save(analysisOrder)
                }

                // 依次进行静态分析，动态分析，模型研判，报告生成
                // 1.静态分析
                logger.info { "$apkName: 开始静态分析" }
                staticAnalysis(analysisOrder, resultPath)

                // 2.动态分析
                GlobalMutex.deviceLock.withLock {
                    logger.info { "$apkName: 开始动态分析" }
                    dynamicAnalysis(analysisOrder, resultPath)
                }

                // 3.研判
                logger.info { "$apkName: 开始进行研判" }
                judge(analysisOrder, resultPath)
                logger.info { "$apkName: 研判完成" }
            } else {
                logger.info { "$apkName 已经分析过" }
                var rejudgeFlag = false  // 若静态或动态分析重新执行，则重新研判
                if (order.staticStatus != "Success" && order.staticStatus != "Analysing") {
                    // 静态分析不成功或正在分析，尝试重新进行静态分析
                    logger.info { "$apkName: 尝试重新进行静态分析" }
                    staticAnalysis(order, resultPath)
                    if (order.staticStatus == "Success") rejudgeFlag = true
                }
                if (order.dynamicStatus != "Success" && order.dynamicStatus != "Analysing" && order.dynamicStatus != "Test Run Failed") {
                    // 动态分析不成功，尝试重新进行动态分析
                    GlobalMutex.deviceLock.withLock {
                        logger.info { "$apkName: 尝试重新进行动态分析" }
                        dynamicAnalysis(order, resultPath)
                        if (order.dynamicStatus == "Success") rejudgeFlag = true
                    }
                }
                if (order.judgeStatus != "Success" || rejudgeFlag) {
                    // 模型研判不成功，尝试重新进行模型研判
                    logger.info { "$apkName 尝试重新进行研判" }
                    judge(order, resultPath)
                    logger.info { "$apkName 研判完成" }
                }
            }
        }
    }

    /**
     * 静态分析
     * @param analysisOrder 当前分析订单
     * @param resultPath 结果文件
     * @return 分析结果，一个 json 字符串
     * */
    private fun staticAnalysis(analysisOrder: AnalysisOrder, resultPath: Path) {
        try {
            val explorer = StaticExplorer(userInputPath.relativeTo(PathUtils.INPUT).resolve(analysisOrder.apkName).toString())
            explorer.dump(resultPath)
            analysisOrder.packageName = explorer.getPackageName()
            analysisOrder.staticStatus = "Success"
            analysisOrderRepository.save(analysisOrder)
        } catch (e: Exception) {
            e.printStackTrace()
            analysisOrder.staticStatus = e.toString()
            analysisOrderRepository.save(analysisOrder)
        }
    }

    /**
     * 动态分析并将分析结果和时间保存到数据库中
     * @param analysisOrder 当前分析订单
     * @param resultPath result.json 结果保存路径
     * */
    private fun dynamicAnalysis(analysisOrder: AnalysisOrder, resultPath: Path) {
        // 获得静态分析结果
        if (analysisOrder.staticStatus != "Success") {
            analysisOrder.dynamicStatus = "No Meta Data"
            analysisOrderRepository.save(analysisOrder)
            return
        }
        val apkInfo = resultPath.readText()
        val dynamicExplorer = DynamicExplorer(resultPath.parent)
        val apk = SerializeUtils.load(apkInfo)

        var startTime = System.currentTimeMillis()
        dynamicExplorer.fromApk(apk)
        val dynamicTime = System.currentTimeMillis() - startTime
        analysisOrder.dynamicTime = dynamicTime

        // 开始记录数据处理阶段耗时
        startTime = System.currentTimeMillis()
        // 检查动态分析状态，并依次处理
        val dynamicStatus = dynamicExplorer.status
        if (!dynamicStatus.existEnvironment) {
            analysisOrder.dynamicStatus = "No Device Found"
        } else if (!dynamicStatus.installStatus) {
            analysisOrder.dynamicStatus = "Install Failed"
        } else if (!dynamicStatus.testRunStatus) {
            analysisOrder.dynamicStatus = "Test Run Failed"
        } else {
            // 进一步处理截图和 pcap 包
            val outputPath = resultPath.parent
            val screenContent = dataProcess.parseScreencap(outputPath, apk.packageName)
            apk.screenContent += screenContent
            for (file in outputPath.listDirectoryEntries()) {
                if (file.extension == "pcap") {
                    val urls = dataProcess.parsePCAP(file)
                    apk.urls.addAll(urls)
                }
            }
            // 保存动态分析结果
            SerializeUtils.dump(apk, resultPath, false)
            logger.info { "APK 信息导出至: ${resultPath.toAbsolutePath()}" }

            // 数据处理完成后再保存记录到数据库中
            analysisOrder.dynamicStatus = "Success"
        }
        val dataProcessTime = System.currentTimeMillis() - startTime
        analysisOrder.dataProcessTime = dataProcessTime
        analysisOrderRepository.save(analysisOrder)
    }

    /**
     * 研判，包括黑名单研判和模型研判，其中模型研判又分为基本信息研判和 app 图像研判
     * @param analysisOrder 当前分析记录
     * @param resultPath result.json 文件的路径
     * */
    fun judge(analysisOrder: AnalysisOrder, resultPath: Path) {
        // 获得静态分析结果
        val startTime = System.currentTimeMillis()
        if (analysisOrder.staticStatus != "Success") {
            // 静态分析失败
            analysisOrder.judgeStatus = "No Meta Data"
            analysisOrder.judgeResult = "-1"
        } else {
            val apkInfo = resultPath.readText()
            val apk = SerializeUtils.load(apkInfo)

            val judgeResult = JudgeResult()

            if (analysisOrder.judgeStatus != "Success") {
                askWhiteList(analysisOrder, apk, judgeResult)
            }
            if (analysisOrder.judgeStatus != "Success") {
                askBlacklist(analysisOrder, apk, judgeResult)
            }
            if (analysisOrder.judgeStatus != "Success") {
                askLLM(analysisOrder, apk, judgeResult)
            }
        }
        val judgeTime = System.currentTimeMillis() - startTime
        analysisOrder.judgeTime = judgeTime
        analysisOrderRepository.save(analysisOrder)
    }

    /**
     * 通过 packageName，md5 等信息进行白名单研判，若研判成功，将 analysisOrder.judgeStatus 设置为 Success
     * @param analysisOrder AnalysisOrder 当前分析记录
     * @param apk APK 当前分析的 apk 信息
     * @param judgeResult JudgeResult 分析结果保存
     */
    private fun askWhiteList(analysisOrder: AnalysisOrder, apk: APK, judgeResult: JudgeResult) {
        if (apk.packageName.isNotBlank()) {
            val possibleApks = apkWhitelistRepository.findAllByPackageName(apk.packageName)
            if (possibleApks.isNotEmpty()) {
                for (possibleApk in possibleApks) {
                    val (matched, possibility) = apk.compareBy(possibleApk.packageName, possibleApk.appName, possibleApk.md5)
                    if (matched && possibility >= 2) {
                        // 白名单匹配成功
                        analysisOrder.judgeStatus = "Success"
                        judgeResult.addReason("白名单匹配成功: [名称：${possibleApk.appName} 包名：${possibleApk.packageName} 摘要：${possibleApk.md5}]", "normal")
                        updateAnalysisOrder(analysisOrder, judgeResult)
                        return
                    }
                }
            }
        }
        if (apk.name.isNotBlank()) {
            val possibleApks = apkWhitelistRepository.findAllByAppName(apk.name)
            if (possibleApks.isNotEmpty()) {
                for (possibleApk in possibleApks) {
                    val (matched, possibility) = apk.compareBy(possibleApk.packageName, possibleApk.appName, possibleApk.md5)
                    if (matched && possibility >= 2) {
                        // 白名单匹配成功
                        analysisOrder.judgeStatus = "Success"
                        judgeResult.addReason("白名单匹配成功: [名称：${possibleApk.appName} 包名：${possibleApk.packageName} 摘要：${possibleApk.md5}]", "normal")
                        updateAnalysisOrder(analysisOrder, judgeResult)
                        return
                    }
                }
            }
        }
        if (apk.getMd5().isNotBlank()) {
            val possibleApks = apkWhitelistRepository.findAllByMd5(apk.getMd5())
            if (possibleApks.isNotEmpty()) {
                // 白名单匹配成功
                analysisOrder.judgeStatus = "Success"
                judgeResult.addReason("白名单匹配成功: [名称：${apk.name} 包名：${apk.packageName} 摘要：${apk.getMd5()}]", "normal")
                updateAnalysisOrder(analysisOrder, judgeResult)
                return
            }
        }
    }

    /**
     * 黑名单研判，若研判成功，将 analysisOrder.judgeStatus 设置为 Success
     * @param analysisOrder 当前分析记录
     * @param apk 当前分析的 apk 信息
     * @param judgeResult 分析结果高精度保存
     * */
    private fun askBlacklist(analysisOrder: AnalysisOrder, apk: APK, judgeResult: JudgeResult) {
        var success = false  // 若匹配到黑名单则设置为 true
        // 可疑签名开发者黑名单匹配
        var subject = apk.v1SignatureInfo.certification.subject
        if (subject == "") subject = apk.v2SignatureInfo.certification.subject
        if (subject == "") subject = apk.v3SignatureInfo.certification.subject
        if (subject == "") logger.error { "${apk.name} 签名信息不存在" }
        signatureBlacklistRepository.findAllBySubject(subject)?.forEach {
            judgeResult.addReason("匹配到可疑开发者签名：${subject}", it.type)
            logger.info { "匹配到可疑开发者签名：${subject}" }
            success = true
        }

        // screen content 黑名单匹配
        val words = dataProcess.wordSeg("${apk.screenContent} ${apk.name}")
        logger.info { "screen content 分词结果为： $words" }
        for (word in words) {
            val wordMatched = wordBlacklistRepository.findByWord(word) ?: continue
            judgeResult.addReason("匹配到可疑文本：${wordMatched.word}", wordMatched.type)
            logger.info { "匹配到可疑文本：${wordMatched.word} (${wordMatched.type})" }
            success = true
        }

        if (success) {
            updateAnalysisOrder(analysisOrder, judgeResult)
        } else {
            analysisOrder.judgeStatus = "Judging"
        }
        analysisOrderRepository.save(analysisOrder)
    }

    /**
     * 模型研判
     * @param analysisOrder 当前分析订单
     * @param apk 当前分析的 apk 信息
     * @param judgeResult 分析结果高精度保存
     * */
    private fun askLLM(analysisOrder: AnalysisOrder, apk: APK, judgeResult: JudgeResult) {
        runBlocking {
            val ask1 = async { gpt.ask(apk) }
            val ask2 = async { gpt.askVision(userOutputPath.resolve(analysisOrder.analysisNo)) }
            val result1 = ask1.await()
            val (type1, reason1) = parseLLMResult(result1)
            val result2 = ask2.await()
            val (type2, reason2) = parseLLMResult(result2)
            when (type1) {
                "赌博" -> judgeResult.addReason(reason1, "gamble")
                "诈骗" -> judgeResult.addReason(reason1, "scam")
                "色情" -> judgeResult.addReason(reason1, "sex")
                "普通" -> judgeResult.addReason(reason1, "normal")
            }
            when (type2) {
                "赌博" -> judgeResult.addReason(reason2, "gamble")
                "诈骗" -> judgeResult.addReason(reason2, "scam")
                "色情" -> judgeResult.addReason(reason2, "sex")
                "普通" -> judgeResult.addReason(reason2, "normal")
            }
            updateAnalysisOrder(analysisOrder, judgeResult)
            logger.info { "${apk.name} 研判结果：${if(judgeResult.isScam) "涉诈" else "不涉诈"}" }
            analysisOrder.judgeStatus = "Success"
            analysisOrderRepository.save(analysisOrder)
        }
    }

    /**
     * 根据分析号找到对应的分析记录
     * */
    fun getAnalysisOrder(analysisNo: String): AnalysisOrder? {
        return analysisOrderRepository.findByAnalysisNo(analysisNo)
    }

    /**
     * 根据分析号得到 APK
     * @param analysisNo 分析号
     * */
    fun getAPK(analysisNo: String): APK {
        val resultPath = userOutputPath.resolve(analysisNo).resolve("result.json")
        if (!resultPath.exists()) {
            return APK()
        }
        return SerializeUtils.load(resultPath.readText())
    }

    /**
     * 获取黑名单列表，也就是返回 analysis_order 的所有数据
     * */
    fun getBlackList(): List<AnalysisOrder> {
        val all = analysisOrderRepository.findAll()
        return all
    }

    /**
     * 处理大模型返回的字符串类型数据
     * @param result 模型研判结果原始数据
     * @return <判断类型，判断理由>
     * */
    private fun parseLLMResult(result: String): Pair<String, String> {
        val lines = result.split("\n").filter { it.trim().isNotEmpty() }  // 用换行符分隔，删除空行
        val reason = lines.drop(1).joinToString(separator = "\n")
        val type = lines[0].lowercase()
        return if (type.contains("yes") || type.contains("赌博") || type.contains("色情") || type.contains("诈骗")) {
            Pair(type, reason)
        } else {
            Pair("普通", reason)
        }
    }

    /**
     * 用 judgeResult 更新 analysisOrder
     * */
    private fun updateAnalysisOrder(analysisOrder: AnalysisOrder, judgeResult: JudgeResult) {
        analysisOrder.judgeStatus = "Success"
        analysisOrder.isGamble = judgeResult.isGamble
        analysisOrder.isScam = judgeResult.isScam
        analysisOrder.isSex = judgeResult.isSex
        analysisOrder.scamReason = judgeResult.scamReason()
        analysisOrder.sexReason = judgeResult.sexReason()
        analysisOrder.gambleReason = judgeResult.gambleReason()
        analysisOrder.judgeResult = judgeResult.normalReason()
    }

    /**
     * 获得所有截图的 uri，例如 /static/0a3a2bfaf53361e1cd2269a26d564848/step-1--g0a0-1720015889675.png
     * @param analysisNo 分析号
     * */
    fun getScreencaps(analysisNo: String): List<String> {
        val targetPath = userOutputPath.resolve(analysisNo)
        if (!targetPath.exists()) {
            return emptyList()
        }
        val ret = mutableSetOf<String>()
        for (file in targetPath.listDirectoryEntries()) {
            if (file.extension == "png") {
                ret.add("/static/$analysisNo/${file.name}")
            }
        }
        return ret.toList()
    }

    /**
     * 根据 order 的 is_XX 和 XX_reason 字段构建字符串列表
     * @param analysisOrder 指定记录
     * */
    fun getJudgeResult(analysisOrder: AnalysisOrder): List<String> {
        var normal = true
        val ret = mutableSetOf<String>()
        if (analysisOrder.isScam) {
            ret.add("诈骗\n${analysisOrder.scamReason}")
            normal = false
        }
        if (analysisOrder.isSex) {
            ret.add("色情\n${analysisOrder.sexReason}")
            normal = false
        }
        if (analysisOrder.isGamble) {
            ret.add("赌博\n${analysisOrder.gambleReason}")
            normal = false
        }
        if (normal) {
            if (analysisOrder.judgeResult.contains("白名单")) {
                ret.add("普通\n${analysisOrder.judgeResult}")
            } else {
                ret.add("普通\n${analysisOrder.judgeResult}")
            }
        }
        return ret.toList()
    }

    /**
     * 根据 analysisNo 删除数据库记录以及 userOutput 下的文件
     * @param analysisNo 分析号
     * */
    @OptIn(ExperimentalPathApi::class)
    fun deleteAnalysis(analysisNo: String) {
        analysisOrderRepository.deleteAnalysisOrderByAnalysisNo(analysisNo)
        val targetPath = userOutputPath.resolve(analysisNo)
        if (targetPath.exists()) {
            targetPath.deleteRecursively()
        }
    }
}
